#include "Common.h"
#include "MeshBuilder.h"
#include "VertexIndexData.h"
#include "Render.h"
#include "Mesh.h"
#include "MeshManager.h"
#include "MaterialManager.h"
#include "Material.h"
#include <d3dx9.h>
#include <set>

void MeshBuilder::Reset()
{
	mBuildingList.clear();
	mIndexList.clear();
	mMesh = NULL;
	mAABB.Init();
}

void MeshBuilder::Begin(const std::string& name, size_t reserve)
{
	mMesh = MeshManager::Instance().Create(name);

	SubMesh* subMesh = mMesh->CreateSubMesh();
	subMesh->setMaterialName(MaterialManager::Instance().getShapeMaterial()->getName());
	subMesh->setShareGeometry(true);

	mBuildingList.reserve(reserve);
	mIndexList.reserve(reserve * 2);
}

void MeshBuilder::AddVertex(
	const Vector3& position,
	const Vector2& texcoord,
	const Vector3& normal,
	const Vector3& tangent)
{
	AddVertex(Vertex(position, normal, tangent, texcoord));
}

void MeshBuilder::AddVertex(const Vertex& vertex)
{
	mAABB.Merge(vertex.Position);

	mBuildingList.push_back(vertex);
}

void MeshBuilder::AddTriangle(unsigned short v0, unsigned short v1, unsigned short v2)
{
	mIndexList.push_back(v0);
	mIndexList.push_back(v1);
	mIndexList.push_back(v2);
}

void MeshBuilder::End()
{
	if (mBuildingList.empty())
		return;

	VertexData::StreamInfo positionVB = BuildPositionVB();
	VertexData::StreamInfo noPositionVB = BuildNoPositionVB();

	IDirect3DIndexBuffer9* indexBuffer = BuildIndexBuffer();
	IndexData::PrimitiveInfo primitiveInfo = {
		D3DPT_TRIANGLELIST,
		0,
		0, mBuildingList.size(),
		0, mIndexList.size() / 3
	};

	VertexData* vertexData = new VertexData(positionVB, noPositionVB);
	IndexData* indexData = new IndexData(indexBuffer, primitiveInfo);

	mMesh->setVertexData(vertexData);
	mMesh->GetSubMesh(0)->setIndexData(indexData);
	mMesh->setBounds(mAABB);
	mMesh->Init();
}

VertexData::StreamInfo MeshBuilder::BuildPositionVB(bool gpuSkin)
{
	struct PackedWeight
	{
		BYTE X, Y, Z, W;
		PackedWeight(const float* weights)
			: X(Clamp(Math::Trunc(weights[0] * 255.0f), 0, 255))
			, Y(Clamp(Math::Trunc(weights[1] * 255.0f), 0, 255))
			, Z(Clamp(Math::Trunc(weights[2] * 255.0f), 0, 255))
			, W(Clamp(Math::Trunc(weights[3] * 255.0f), 0, 255))
		{ }
	};

	struct PackedVertex
	{
		Vector3 Position;
		DWORD BoneIndies;
		PackedWeight BoneWeight;

		PackedVertex(const Vertex& vertex)
			: Position(vertex.Position), BoneIndies(*(DWORD*)vertex.InfluenceBones)
			, BoneWeight(vertex.InfluenceWeights)
		{ }
	};

	size_t vertexSize = gpuSkin ? sizeof(PackedVertex) : sizeof(Vector3);
	size_t vertexBufferSize = vertexSize * mBuildingList.size();
	
	// create the vertex buffer.
	IDirect3DVertexBuffer9* vertexBuffer = CreateVB(vertexBufferSize, 0, D3DPOOL_MANAGED);

	// get the memory pointer
	void* vbData = NULL;
	if (FAILED(vertexBuffer->Lock(0, 0, &vbData, 0)))
	{
		SAFE_RELEASE(vertexBuffer);
		return VertexData::StreamInfo();
	}

	// fill data one by one.
	Vector3* positionData = (Vector3*)vbData;
	PackedVertex* vertexData = (PackedVertex*)vbData;
	for (VertexList::iterator it(mBuildingList.begin()), end(mBuildingList.end()); it != end; it++)
	{
		if (gpuSkin)
			*vertexData++ = PackedVertex(*it);
		else
			*positionData++ = it->Position;
	}

	vertexBuffer->Unlock();

	return VertexData::StreamInfo(0, vertexBuffer, 0, vertexSize);
}

VertexData::StreamInfo MeshBuilder::BuildNoPositionVB()
{
	struct PackedVertex
	{
		PackedVector Normal;
		PackedVector Tangent;
		D3DXVECTOR2_16F Texcoord;

		PackedVertex(const Vertex& vertex)
			: Normal(vertex.Normal), Tangent(vertex.Tangent), Texcoord(vertex.Texcoord.X, vertex.Texcoord.Y)
		{ }
	};

	size_t vertexSize = sizeof(PackedVertex);
	size_t vertexBufferSize = vertexSize * mBuildingList.size();
	
	// create the vertex buffer.
	IDirect3DVertexBuffer9* vertexBuffer = CreateVB(vertexBufferSize, 0, D3DPOOL_MANAGED);

	// get the memory pointer
	void* vbData = NULL;
	if (FAILED(vertexBuffer->Lock(0, 0, &vbData, 0)))
	{
		SAFE_RELEASE(vertexBuffer);
		return VertexData::StreamInfo();
	}

	// fill data one by one.
	PackedVertex* vertexData = (PackedVertex*)vbData;
	for (VertexList::iterator it(mBuildingList.begin()), end(mBuildingList.end()); it != end; it++)
	{
		PackedVertex vertex(*it);
		*vertexData++ = vertex;
	}

	vertexBuffer->Unlock();
	
	return VertexData::StreamInfo(1, vertexBuffer, 0, vertexSize);
}

IDirect3DIndexBuffer9* MeshBuilder::BuildIndexBuffer()
{
	if (mIndexList.empty())
		return NULL;

	size_t indexBufferSize = sizeof(unsigned short) * mIndexList.size();

	// create the index buffer.
	IDirect3DIndexBuffer9* indexBuffer = CreateIB(indexBufferSize, 0, D3DPOOL_MANAGED);

	// get the memory pointer
	void* ibData = NULL;
	if (FAILED(indexBuffer->Lock(0, 0, &ibData, 0)))
	{
		SAFE_RELEASE(indexBuffer);
		return NULL;
	}

	// fill data one by one.
	unsigned short* indexData = (unsigned short*)ibData;
	for (IndexList::iterator it(mIndexList.begin()), end(mIndexList.end()); it != end; it++)
		*indexData++ = *it;

	indexBuffer->Unlock();

	return indexBuffer;
}

IDirect3DVertexBuffer9* MeshBuilder::CreateVB(size_t size, size_t usage, size_t pool)
{
	LPDIRECT3DDEVICE9 device = Render::Instance().getDevice();
	IDirect3DVertexBuffer9* ret = NULL;
	HRESULT hr = device->CreateVertexBuffer(size, usage, 0, (D3DPOOL)pool, &ret, NULL);
	return ret;
}

IDirect3DIndexBuffer9* MeshBuilder::CreateIB(size_t size, size_t usage, size_t pool)
{
	LPDIRECT3DDEVICE9 device = Render::Instance().getDevice();
	IDirect3DIndexBuffer9* ret = NULL;
	HRESULT hr = device->CreateIndexBuffer(size, usage, D3DFMT_INDEX16, (D3DPOOL)pool, &ret, NULL);
	return ret;
}

bool MeshBuilder::BuildMesh(MeshData& meshData, Mesh* mesh)
{
	Reset();

	// collect all of the verties.
	ProcessVertexBuffer(meshData.SharedVertexBuff, meshData.VertexBoneAssignments, mesh->getBoneIndexMap());
	for (SubMeshDataList::const_iterator it(meshData.SubMeshes.begin()), end(meshData.SubMeshes.end()); it != end; it++)
	{
		SubMesh* subMesh = mesh->CreateSubMesh();
		subMesh->setName(it->Name);
		subMesh->setMaterialName(it->Material);
		subMesh->setShareGeometry(true);

		size_t baseVertexCount = mBuildingList.size();
		ProcessVertexBuffer(it->PrivateVertexBuff, it->VertexBoneAssignments, subMesh->getBoneIndexMap());
		
		// clear the index data.
		mIndexList.clear();
		mIndexList.reserve(it->IndexBuff.IndexCount);

		const WORD* ibuff16 = (const WORD*)&it->IndexBuff.Buff[0];
		const DWORD* ibuff32 = (const DWORD*)&it->IndexBuff.Buff[0];
		for (size_t i = 0; i < it->IndexBuff.IndexCount; i++, ibuff16++, ibuff32++)
		{
			DWORD index = it->IndexBuff.Uint32Index ? *ibuff32 : *ibuff16;
			mIndexList.push_back(static_cast<WORD>(index));
		}

		// build the index buffer.
		IDirect3DIndexBuffer9* indexBuffer = BuildIndexBuffer();
		IndexData::PrimitiveInfo primitiveInfo = {
			D3DPT_TRIANGLELIST,
			it->UseSharedGeometry ? 0 : baseVertexCount, // shared vertex start at 0.
			0, 
			it->UseSharedGeometry ? mBuildingList.size() : it->PrivateVertexBuff.VertexCount,
			0, 
			mIndexList.size() / 3
		};

		IndexData* indexData = new IndexData(indexBuffer, primitiveInfo);
		subMesh->setIndexData(indexData);
	}

	// build the vertex buffer.
	VertexData::StreamInfo positionVB = BuildPositionVB(meshData.SkeletallyAnimated);
	VertexData::StreamInfo noPositionVB = BuildNoPositionVB();
	VertexData* vertexData = new VertexData(positionVB, noPositionVB);
	mesh->setVertexData(vertexData);

	// setup in to the bounds.
	if (meshData.Bounds.Valid())
	{
		mesh->setBounds(meshData.Bounds);
		mesh->setBoundRadius(meshData.BoundRadius);
	}
	else
		mesh->setBounds(mAABB);

	return true;
}

bool MeshBuilder::ProcessVertexBuffer(
	const VertexDataBuff& vb, 
	const VertexBoneAssignmentList& vbas, 
	IndexMap& indexMap)
{
	// check the vb is valid.
	if (vb.VertexCount == 0 || vb.VertexSize == 0 || vb.VertexDeclares.empty() || vb.Buff.empty())
		return false;

	typedef const VertexDeclare* PVertexDeclare;
	PVertexDeclare posDesc = NULL, normalDesc = NULL, tangentDesc = NULL, uvDesc = NULL;
	for (VertexDeclareList::const_iterator it(vb.VertexDeclares.begin()), end(vb.VertexDeclares.end()); it != end; it++)
	{
		switch (it->semantic)
		{
		case VES_POSITION:
			posDesc = &(*it);
			break;
		case VES_NORMAL:
			normalDesc = &(*it);
			break;
		case VES_TANGENT:
			tangentDesc = &(*it);
			break;
		case VES_TEXTURE_COORDINATES:
			uvDesc = &(*it);
			break;
		}
	}
	if (posDesc == NULL)
		return false; // no position buffer.
	
	size_t oldSize = mBuildingList.size();

	// fetch the vertex data.
	mBuildingList.reserve(mBuildingList.size() + vb.VertexCount);
	for (DWORD i = 0; i < vb.VertexCount; i++)
	{
		Vertex vertex;
		vertex.Position = *(const Vector3*)(&vb.Buff[i * vb.VertexSize + posDesc->offset]);

		mAABB.Merge(vertex.Position);

		if (normalDesc != NULL)
		{
			vertex.Normal = *(const Vector3*)(&vb.Buff[i * vb.VertexSize + normalDesc->offset]);
			vertex.Normal.Normalize();
		}

		if (tangentDesc != NULL)
			vertex.Tangent = *(const Vector3*)(&vb.Buff[i * vb.VertexSize + tangentDesc->offset]);

		if (uvDesc != NULL)
			vertex.Texcoord = *(const Vector2*)(&vb.Buff[i * vb.VertexSize + uvDesc->offset]);

		// reset the bone influence to zero.
		for (int j = 0; j < MAX_INFLUENCES; j++) {
			vertex.InfluenceBones[j] = 0;
			vertex.InfluenceWeights[j] = 0;
		}

		mBuildingList.push_back(vertex);
	}
	size_t newSize = mBuildingList.size();
	size_t vertexCount = newSize - oldSize;

	if (!vbas.empty())
	{
		// build bone inflence.
		typedef std::set<WORD> BoneIndexSet;
		BoneIndexSet usedBoneIndices;

		for (VertexBoneAssignmentList::const_iterator it(vbas.begin()), end(vbas.end()); it != end; it++)
			usedBoneIndices.insert(it->second.BoneIndex);
		
		size_t index = 0;
		IndexMap blendIndexMap;

		indexMap.resize(usedBoneIndices.size());
		blendIndexMap.resize(*usedBoneIndices.rbegin() + 1);

		for (BoneIndexSet::const_iterator it(usedBoneIndices.begin()), end(usedBoneIndices.end()); it != end; it++, index++)
		{
			indexMap[index] = *it;
			blendIndexMap[*it] = index;
		}
		
		VertexBoneAssignmentList::const_iterator vbaIt(vbas.begin()), vbaEnd(vbas.end());
		for (index = 0; index < vertexCount; index++)
		{
			Vertex& vertex = mBuildingList[index + oldSize];
			for (int i = 0; i < MAX_INFLUENCES; i++)
			{
				if (vbaIt != vbaEnd && vbaIt->second.VertexIndex == index)
				{
					vertex.InfluenceBones[i] = blendIndexMap[vbaIt->second.BoneIndex];
					vertex.InfluenceWeights[i] = vbaIt->second.Weight;
					vbaIt++;
				}
				else
				{
					vertex.InfluenceBones[i] = 0;
					vertex.InfluenceWeights[i] = 0;
				}
			}

			// if there is more vba exist.
			while (vbaIt != vbaEnd && vbaIt->second.VertexIndex == index)
				vbaIt++;
		}
	}

	return true;
}
